# -*- coding: utf-8 -*-
"""
Created on Sun Oct 11 10:09:35 2020
用于双量子点量子模型
单量子比特的归零操作
多个初始态对多个目标态的归零保真度

@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy

np.random.seed(1)
T = np.pi
dt = np.pi/20
step_max = T/dt
sx = np.mat([[0, 1], [1, 0]], dtype=complex)
sy = np.mat([[0, -1j], [1j, 0]], dtype=complex)
sz = np.mat([[1, 0], [0, -1]], dtype=complex)

# a0,a1,a2,a3,a4,a5,a6,a7,a8,a9 = 0,0,0,0,0,0,0,0,0,0 #统计各动作被选用的频率

action_space=np.array([[1,0,0], #可以选择的动作范围，各列的每项分别代表着 sigma x, y, z 前面的系数。
                       [2,0,0], #每次执行的动作都是单独的绕 x, y, z 轴一定角度的旋转
                       [0,1,0], # x, y 方向的值可以取负，但 z 方向的只能取正值
                       [0,2,0],
                       [0,0,1],
                       [0,0,2],
                       [-1,0,0],
                       [-2,0,0],
                       [0,-1,0],
                       [0,-2,0]])

noise_normal = np.array([ 1.62434536, -0.61175641, -0.52817175, -1.07296862,  0.86540763,
                         -2.3015387 ,  1.74481176, -0.7612069 ,  0.3190391 , -0.24937038,
                          1.46210794, -2.06014071, -0.3224172 , -0.38405435,  1.13376944,
                         -1.09989127, -0.17242821, -0.87785842,  0.04221375,  0.58281521,
                         -1.10061918,  1.14472371,  0.90159072,  0.50249434,  0.90085595,
                         -0.68372786, -0.12289023, -0.93576943, -0.26788808,  0.53035547,
                         -0.69166075, -0.39675353, -0.6871727 , -0.84520564, -0.67124613,
                         -0.0126646 , -1.11731035,  0.2344157 ,  1.65980218,  0.74204416,
                         -0.19183555])
#noise_normal 为均值为 0 ，标准差为 1 的正态分布的随机数组成的数组
#该随机数由 np.random.seed(1) 生成: np.random.seed(1) \ noise_uniform = np.random.normal(loc=0.0, scale=1.0, size=41)

noise_drift = np.array([1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,])


theta_num = 6 #除了 0 和 Pi 两个点之外，点的数量
varphi_num = 21#varphi 角度一圈上的点数

theta = np.linspace(0,np.pi,theta_num+2,endpoint=True) 
varphi = np.linspace(0,np.pi*2,varphi_num,endpoint=False) 

def psi_set():
    psi_set = []
    for ii in range(1,theta_num+1):
        for jj in range(varphi_num):
            psi_set.append(np.mat([[np.cos(theta[ii]/2)],[np.sin(theta[ii]/2)*(np.cos(varphi[jj])+np.sin(varphi[jj])*(0+1j))]]))
    psi_set.append(np.mat([[1], [0]], dtype=complex))
    psi_set.append(np.mat([[0], [1]], dtype=complex))
    return psi_set

target_set = psi_set()
init_set = psi_set()

#----------------------------------------------------------
def step(psi,target_psi,action):
    
    H = action[0]*sx/2 + action[1]*sy/2 - action[2]*sz/2
     
    U = expm(-1j * H * dt) 
    psi_ = U * psi
    fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real
    return fid, psi_

#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        psi_list.append(psi_)
        fid_list.append(fid)
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
    
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    action_list = list(range(len(action_space)))
    for action in action_list:
        
        H = float(action_space[action,0])*sx/2 + float(action_space[action,1])*sy/2 - float(action_space[action,2])*sz/2
        U = expm(-1j * H * dt) 
        psi_ = U * psi
        fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
        
        psi_list.append(psi_)
        fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
        # best_action = fid_list.index(min(fid_list))
        # best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_

def job(target_psi):
    
    fid_list_noise = [] # 中间元素为列表，每个列表元素为[a,b,c,d,...],其中 a 为无噪声时， 其他为有噪声时
    start_time = time()
    
    for psi1 in init_set:
        
        fid_list = [] # 元素为[a,b,c,d,...],其中 a 为无噪声时， 其他为有噪声时
        
        action_list_1 = []
        action_list_2 = []
        action_list_0 = []
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        
        fid_max = F
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        fid_list_1 = [F]
        fid_list_2 = [F]
        fid_list_0 = [F]
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            action_list_1.append(action)
            fid_list_1.append(F)            
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            action_list_2.append(action)
            fid_list_2.append(F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            action_list_0.append(action)
            fid_list_0.append(F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max0,fid_max1,fid_max2) 
        fid_list.append(fid_max)
        fid_max_index = [fid_max1,fid_max2,fid_max0].index(max([fid_max1,fid_max2,fid_max0]))
        
        if fid_max_index == 0: 
            actions_list = action_list_1
            fids_list = fid_list_1
        elif fid_max_index == 1:
            actions_list = action_list_2
            fids_list = fid_list_2
        elif fid_max_index == 2:
            actions_list = action_list_0
            fids_list = fid_list_0
        
        max_index = fids_list.index(fid_max)
        actions_list = actions_list[0:max_index+1]
        
        #在对应于最大保真度的动作策略上施加噪声
         
        for noise_am in noise_am_list:
            
            noise_list = noise_am * noise_drift 
            # noise_list = noise_am * noise_normal
            psi = psi1
            F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
            fid_noise_list = [F]
            i = 0
            for action_index in actions_list:
                action = action_space[action_index]*(1+noise_list[i])
                i += 1
                F, psi_ =  step(psi, target_psi, action)
                fid_noise_list.append(F)
                psi = psi_
            fid_noise = fid_noise_list[max_index]
            fid_list.append(fid_noise)
            
        fid_list_noise.append(fid_list)
        
    fid_list_noise = np.array(fid_list_noise)
            
    return np.mean(fid_list_noise,axis=0)
        
    # return  np.mean(fids_list)

def multicore():o
    pool = mp.Pool()
    F_list_noise = pool.map(job, target_set)
    
    return np.mean(np.array(F_list_noise),axis=0)

# noise_am_list = [0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10,0.20,0.30,0.40,0.50]
noise_am_list = [-0.10,-0.09,-0.08,-0.07,-0.06,-0.05,-0.04,-0.03,-0.02,-0.01,0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10]
    

if __name__ == '__main__':
    # print(target_set)
    time1 = time()
    
    F_list_noise =  multicore()

    print('\n F_list_noise = \n',F_list_noise) 
    
    time2 = time()
    
    print('\n time_cost is: ',time2-time1)

# 128测试点
#动作 x,y: 0,1,2,-1,-2; z: 0,1,2

# noise normal 
# noise_am_list = [0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10,0.20,0.30,0.40,0.50]

# F_list_noise = 
#   [0.99944394 0.99944394 0.99942998 0.99939068 0.99932605 0.9992361 0.99912085 0.99898033 0.99881457 0.99862359 0.99840745 0.99816617
#   0.99438406 0.9881705  0.97964254 0.96895775]

# #  time_cost is:  105.54917049407959

# noise drift
# noise_am_list = [-0.10,-0.09,-0.08,-0.07,-0.06,-0.05,-0.04,-0.03,-0.02,-0.01,0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10]

# F_list_noise = 
#  [0.99944394 0.99168881 0.99315973 0.99447683 0.99563969 0.99664795 0.99750138 0.99819981 0.9987432  0.99913157 0.99936508 0.99944394 0.99936848 0.99913912 0.99875638 0.99822085 0.99753324 0.99669432 0.99570499 0.99456621 0.99327902 0.99184458]

#  time_cost is:  113.77585768699646

